#include "net.hpp"

Net::Net() {}

Net::~Net() {}

Net::Net(const int numSeg, const std::vector<float> rangeX_)
{
    numSegement = numSeg;
    rangeX[0]   = rangeX_[0];
    rangeX[1]   = rangeX_[1];

    nearestIdxs.resize(2 * (numSeg - 1));
    // 初始均匀划分
    torch::Tensor temp    = torch::zeros({numSeg - 1, 1});
    const float segLength = (rangeX_[1] - rangeX_[0]) / numSeg;
    for (int i = 0; i < numSeg - 1; ++i)
    {
        temp[i][0] = rangeX_[0] + (i + 1) * segLength;
    }

    segPoi = register_parameter("segPoi", temp, true);
    fit    = register_module("fit", torch::nn::Linear(1, numSeg));
}

at::Tensor Net::forward(at::Tensor x)
{
    const int numInp = x.size(0);

    // 全连接神经网络向前传播，输入为数据点和分段点
    at::Tensor ys = fit->forward(torch::cat({x, segPoi}, 0));

    // 通过分段点信息筛选构造输出
    at::Tensor y_ = torch::zeros({numInp + numSegement - 1, 1});
    for (int i = 0; i < numInp; ++i)
    { // 针对数据点
        int idx = numSegement - 1;
        for (int k = 0; k < numSegement - 1; ++k)
        {
            if (x[i][0].item<float>() < segPoi[k][0].item<float>())
            {
                idx = k;
                break;
            }
        }
        y_[i][0] = ys[i][idx];
    }
    for (int i = numInp; i < numInp + numSegement - 1; ++i)
    { // 针对分段点
        y_[i][0] = ys[i][i - numInp + 1] - ys[i][i - numInp];
    }
    return y_;
}

void Net::printSegementPoint()
{
    for (int i = 0; i < numSegement - 1; ++i)
    {
        std::cout << segPoi[i][0].item<float>() << " ";
    }
    std::cout << std::endl;
}

void Net::checkSegementPoint()
{
    torch::Tensor temp    = torch::zeros({numSegement - 1, 1});
    const float segLength = (rangeX[1] - rangeX[0]) / numSegement;
    for (int i = 0; i < numSegement - 1; ++i)
    { // 默认重置为均匀分段
        temp[i][0] = rangeX[0] + (i + 1) * segLength;
    }
    for (int i = 0; i < numSegement - 1; ++i)
    { // 不能小于定义域最小值
        if (segPoi[i][0].item<float>() < rangeX[0])
        {
            segPoi.set_data(temp);
            return;
        }
    }
    for (int i = 0; i < numSegement - 1; ++i)
    { // 不能大于定义域最大值
        if (segPoi[i][0].item<float>() > rangeX[1])
        {
            segPoi.set_data(temp);
            return;
        }
    }
    for (int i = 0; i < numSegement - 2; ++i)
    { // 分段点不能有交叉
        if (segPoi[i][0].item<float>() > segPoi[i + 1][0].item<float>())
        {
            segPoi.set_data(temp);
            return;
        }
    }
}

CustomDtaset::CustomDtaset(at::Tensor input_data, at::Tensor label_data)
{
    inputs_ = input_data;
    labels_ = label_data;
}

torch::data::Example<> CustomDtaset::get(size_t index)
{
    torch::Tensor Input = inputs_[index];
    torch::Tensor Label = labels_[index];
    return {Input.clone(), Label.clone()};
}

torch::optional<size_t> CustomDtaset::size() const { return labels_.size(0); }
